package com.atbeijing.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.{Encoder, Encoders, SparkSession}

object SparkSQL07_UDAF_Class_1 {

    def main(args: Array[String]): Unit = {

        // TODO 创建环境对象
        val sparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL")
        val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
        import spark.implicits._

        // RDD
        val rdd = spark.sparkContext.makeRDD(
            List(
                (1, "zhangsan", 30),
                (2, "lisi", 40),
                (3, "wangwu", 50)
            )
        )

        val df = rdd.toDF("id", "name", "age")

        df.createOrReplaceTempView("user")

        val udaf = new MyAgeAvgUDAF()
        // 早期版本的Spark无法将强类型的聚合函数应用在SQL文中。
        // 强类型聚合函数在早期版本中使用DSL方式来操作
        val ds = df.as[User]

        // 强类型聚合函数可以转换为查询列进行操作
        ds.select(udaf.toColumn).show


        spark.stop()
    }
    //中间缓存值
    case class AgeBuffer( var total:Int, var cnt:Int )
    //输入值
    case class User(id:Int, name:String, age:Int)
    // TODO 自定义用户聚合函数（UDAF）- 强类型
    // 1. 继承Aggregator
    // 2. 定义泛型
    //    IN  : User
    //    BUF : AgeBuffer
    //    OUT : Int
    // 3. 重写方法
    class MyAgeAvgUDAF extends Aggregator[User, AgeBuffer, Int]{
        // 缓冲区的初始化
        override def zero: AgeBuffer = {
            AgeBuffer(0, 0)
        }

        // 根据输入的值更新缓冲区
        override def reduce(buf: AgeBuffer, user: User): AgeBuffer = {
            buf.total += user.age
            buf.cnt += 1
            buf
        }

        // 缓冲区的合并
        override def merge(b1: AgeBuffer, b2: AgeBuffer): AgeBuffer = {
            b1.total += b2.total
            b1.cnt += b2.cnt
            b1
        }

        // 计算
        override def finish(buf: AgeBuffer): Int = {
            buf.total / buf.cnt
        }

        override def bufferEncoder: Encoder[AgeBuffer] = Encoders.product

        override def outputEncoder: Encoder[Int] = Encoders.scalaInt
    }
}
